package com.wust.dns_resolver.utils;

import lombok.extern.slf4j.Slf4j;

import java.io.*;
import java.net.*;
import java.util.Random;

@Slf4j
public class DnsUtils {

    private static final String DNS_SERVER = "8.8.8.8";//谷歌DNS服务器
    private static final int DNS_PORT = 53;//DNS服务器端口

    /**
     *验证域名格式(前端已验证，此处做后端验证防止抓包传不符要求的数据)
     * @param domainName 域名
     * @return true or false
     */
    public static boolean isValidDomain(String domainName) {

        if (domainName == null || domainName.isEmpty()) return false;
        //域名格式验证
        String regex = "^([a-zA-Z0-9]([a-zA-Z0-9\\-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,}$";

        return domainName.matches(regex);
    }

    /**
     * 构建DNS查询报文
     * @param domainName 域名
     * @param isIPv6 是否查询IPv6地址
     * @return DNS查询报文(字节数组)
     */
    public static byte[] buildQuery(String domainName, boolean isIPv6) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);

        try {
            // 生成随机事务ID
            Random random = new Random();
            int transactionId = random.nextInt(65535);
            dos.writeShort(transactionId);

            // 标志位 (标准查询)
            dos.writeShort(0x0100);

            // 问题数量
            dos.writeShort(1);
            // 资源记录数量
            dos.writeShort(0);
            // 授权资源记录数量
            dos.writeShort(0);
            // 额外资源记录数量
            dos.writeShort(0);

            // 编码域名
            String[] parts = domainName.split("\\.");
            for (String part : parts) {
                dos.writeByte(part.length());
                dos.writeBytes(part);
            }
            dos.writeByte(0); // 域名结束符

            // 查询类型 (A记录或AAAA记录)
            dos.writeShort(isIPv6 ? 28 : 1);  // 28是AAAA记录，1是A记录
            // 查询类 (IN)
            dos.writeShort(1);

        } catch (IOException e) {
            e.printStackTrace();
        }

        return baos.toByteArray();
    }

    /**
     * 打印报文内容
     * @param packet DNS报文(字节数组)
     */
    public static void printPacket(byte[] packet) {

        // 打印报文头部
        log.info("报文头部：");
        log.info("事务ID: 0x{}", String.format("%04X", ((packet[0] & 0xFF) << 8) | (packet[1] & 0xFF)));
        log.info("标志位: 0x{}", String.format("%04X", ((packet[2] & 0xFF) << 8) | (packet[3] & 0xFF)));
        log.info("问题数: {}", ((packet[4] & 0xFF) << 8) | (packet[5] & 0xFF));
        log.info("回答数: {}", ((packet[6] & 0xFF) << 8) | (packet[7] & 0xFF));
        log.info("授权数: {}", ((packet[8] & 0xFF) << 8) | (packet[9] & 0xFF));
        log.info("附加数: {}", ((packet[10] & 0xFF) << 8) | (packet[11] & 0xFF));

        // 打印十六进制dump
        log.info("\n完整报文内容（十六进制）：");

        StringBuilder hexDump = new StringBuilder();

        for (int i = 0; i < packet.length; i++) {
            if (i % 16 == 0) {
                if (i != 0) {
                    log.info(hexDump.toString());
                    hexDump.setLength(0); // 清空StringBuilder
                }
                hexDump.append(String.format("%04X: ", i));
            }
            hexDump.append(String.format("%02X ", packet[i] & 0xFF));
        }

        if (hexDump.length() > 0) {
            log.info(hexDump.toString()); // 打印最后一行
        }
        log.info("\n");
    }

    /**
     * 发送查询并接收响应
     * @param query DNS查询报文
     * @return DNS响应报文(字节数组)
     * @throws IOException 异常
     */
    public static byte[] sendAndReceive(byte[] query) throws IOException {

        DatagramSocket socket = new DatagramSocket();
        socket.setSoTimeout(5000); // 设置超时时间为5秒

        // 发送查询
        InetAddress dnsServer = InetAddress.getByName(DNS_SERVER);
        DatagramPacket sendPacket = new DatagramPacket(query, query.length, dnsServer, DNS_PORT);
        socket.send(sendPacket);

        // 接收响应
        byte[] receiveData = new byte[1024];
        DatagramPacket receivePacket = new DatagramPacket(receiveData, receiveData.length);
        socket.receive(receivePacket);

        socket.close();

        // 返回实际接收到的数据
        byte[] actualData = new byte[receivePacket.getLength()];
        System.arraycopy(receivePacket.getData(), 0, actualData, 0, receivePacket.getLength());
        return actualData;
    }

    /**
     * 解析DNS响应
     * @param response DNS响应报文(字节数组)
     * @param isIPv6 是否解析IPv6地址
     * @return 解析到的IP地址，如果没有找到则返回null
     */
    public static String parseResponse(byte[] response, boolean isIPv6) {
        DnsResponseParser parser = new DnsResponseParser(response);
        String ipAddress = null;

        try {
            // 跳过报文头部（12字节）
            parser.skip(12);

            // 跳过查询部分
            String queryDomain = parser.readDomainName();
            log.info("查询的域名: {}", queryDomain);
            
            parser.skip(4); // 跳过查询类型和查询类

            // 解析应答部分
            boolean found = false;
            int answers = (response[6] << 8) & 0xFF00 | response[7] & 0xFF;

            log.info("\n查询结果：");
            for (int i = 0; i < answers; i++) {
                String domain = parser.readDomainName();
                int type = parser.readShort();
                int dnsClass = parser.readShort();
                log.info("dnsClass:{}",dnsClass);
                long ttl = parser.readInt() & 0xFFFFFFFFL;
                int dataLength = parser.readShort();

                if ((isIPv6 && type == 28) || (!isIPv6 && type == 1)) { // AAAA记录或A记录
                    found = true;
                    byte[] addr = parser.readBytes(dataLength);
                    
                    if (isIPv6) {
                        // 构建IPv6地址字符串
                        StringBuilder sb = new StringBuilder();
                        for (int j = 0; j < 16; j += 2) {
                            if (j > 0) sb.append(":");
                            sb.append(String.format("%02x%02x", addr[j] & 0xFF, addr[j + 1] & 0xFF));
                        }
                        ipAddress = sb.toString();
                    } else {
                        // IPv4地址格式化
                        ipAddress = String.format("%d.%d.%d.%d",
                                addr[0] & 0xFF,
                                addr[1] & 0xFF,
                                addr[2] & 0xFF,
                                addr[3] & 0xFF);
                    }

                    log.info("域名: {}", domain);
                    log.info("记录类型: {}", type);
                    log.info("TTL: {}秒", ttl);
                    log.info("IP地址: {}", ipAddress);
                } else {
                    parser.skip(dataLength);
                }
            }

            if (!found) {
                log.info("未找到对应的A记录");
                return "未找到对应的A记录";
            }

        } catch (Exception e) {
            log.error("解析响应时发生错误: {}", e.getMessage());
            return "解析响应时发生错误";
        }

        return ipAddress;
    }

}
